#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import asyncio
import json
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import List, Dict, Optional, Union

from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from starlette.routing import Route, Mount
from strenum import StrEnum

import imaplib
import smtplib
import email
from email.header import decode_header, make_header
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart

import os
import mcp.types as types
from mcp.server.lowlevel import Server
from mcp.server.sse import SseServerTransport


# -----------------------------
# Configuration & Setup
# -----------------------------

EMAIL_CONFIG = {
    "imap_server": "imap.163.com",
    "imap_port": 993,
    "smtp_server": "smtp.163.com",
    "smtp_port": 465,
    "email": os.getenv("MCP_EMAIL", "your_email@163.com"),
    "password": os.getenv("MCP_EMAIL_PASSWORD", "your_authorization_code")
}

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# -----------------------------
# Email Client Logic
# -----------------------------

class EmailClient:
    def __init__(self, config):
        self.config = config
        self.imap = None
        self.smtp = None
        self.current_folder = None  
        self.connected = False      
        self.lock = asyncio.Lock()

    async def connect(self):
        loop = asyncio.get_event_loop()
        self.imap = await loop.run_in_executor(None, imaplib.IMAP4_SSL, self.config["imap_server"], self.config["imap_port"])
        await loop.run_in_executor(None, self.imap.login, self.config["email"], self.config["password"])

        self.smtp = await loop.run_in_executor(None, smtplib.SMTP_SSL, self.config["smtp_server"], self.config["smtp_port"])
        await loop.run_in_executor(None, self.smtp.login, self.config["email"], self.config["password"])
        
        self.connected = True  
        # await self._ensure_folder_selected("INBOX")  
        # await self.list_mailboxes()    
        logger.info("Connected to email servers.")

    async def disconnect(self):
        if not self.connected:      
            return
        
        if self.imap:
            await asyncio.get_event_loop().run_in_executor(None, self.imap.logout)
        if self.smtp:
            await asyncio.get_event_loop().run_in_executor(None, self.smtp.quit)

        self.connected = False      
        self.current_folder = None  
        logger.info("Disconnected from email servers.")

    async def _ensure_folder_selected(self, folder):  
        """确保已选择指定文件夹，如果未选择则选择它"""
        if not self.connected:
            raise Exception("Not connected to email server")
            
        if self.current_folder != folder:
            try:
                loop = asyncio.get_event_loop()
                status, data = await loop.run_in_executor(None, self.imap.select, folder)
                if status != 'OK':
                    logger.error(f"Failed to select folder {folder}: {data}")
                    self.current_folder = None
                    raise Exception(f"IMAP select failed: {status}, {data}")
                self.current_folder = folder
                logger.info(f"Selected folder: {folder}")
            except Exception as e:
                logger.error(f"Failed to select folder {folder}: {str(e)}")
                self.current_folder = None
                raise

    async def fetch_emails(self, folder="INBOX", limit=10, search_criteria="ALL"):
        async with self.lock:
            try:
                # 确保文件夹已选择
                await self._ensure_folder_selected(folder)  
                
                loop = asyncio.get_event_loop()
                # await loop.run_in_executor(None, self.imap.select, folder)
                status, messages = await loop.run_in_executor(None, self.imap.search, None, search_criteria)
                message_ids = messages[0].split()[-limit:]
                
                if status != 'OK':                          
                        raise Exception(f"Search failed: {status}")
                
                emails = []
                for msg_id in reversed(message_ids):
                    status, msg_data = await loop.run_in_executor(None, self.imap.fetch, msg_id, "(RFC822)")
                    raw_email = msg_data[0][1]
                    msg = email.message_from_bytes(raw_email)

                    if status != 'OK':                          
                            logger.warning(f"Failed to fetch email {msg_id}")
                            continue

                    subject = str(make_header(decode_header(msg["Subject"])))
                    from_ = str(make_header(decode_header(msg["From"])))
                    to = str(make_header(decode_header(msg.get("To", ""))))
                    date = msg.get("Date")
                    body = ""

                    if msg.is_multipart():
                        for part in msg.walk():
                            content_type = part.get_content_type()
                            content_disposition = str(part.get("Content-Disposition"))
                            try:
                                body_part = part.get_payload(decode=True).decode()
                            except:
                                continue
                            if content_type == "text/plain" and "attachment" not in content_disposition:
                                body = body_part
                                break
                    else:
                        body = msg.get_payload(decode=True).decode()

                    emails.append({
                        "subject": subject,
                        "from": from_,
                        "to": to,
                        "date": date,
                        "body": body,
                        "id": msg.get("Message-ID")
                    })
                return emails
            
            except Exception as e:                      
                logger.error(f"Error fetching emails: {str(e)}")
                # 发生错误时重置状态
                self.current_folder = None              
                raise

    async def reply_to_email(self, email_id: str, to: str, body: str):
        msg = MIMEMultipart()
        msg["From"] = self.config["email"]
        msg["To"] = to
        msg["In-Reply-To"] = email_id
        msg["Subject"] = f"Re: {await self._get_subject_by_id(email_id)}"
        msg.attach(MIMEText(body, "plain", "utf-8"))
        await asyncio.get_event_loop().run_in_executor(None, self.smtp.send_message, msg)

    async def forward_email(self, email_id: str, to: str, body: str = "Forwarded message"):
        raw_email = await self._fetch_raw_email_by_id(email_id)
        forwarded_msg = MIMEMultipart()
        forwarded_msg["From"] = self.config["email"]
        forwarded_msg["To"] = to
        forwarded_msg["Subject"] = f"Fwd: {await self._get_subject_by_id(email_id)}"
        forwarded_msg.attach(MIMEText(body + "\n\n--- Original Message ---\n" + raw_email, "plain", "utf-8"))
        await asyncio.get_event_loop().run_in_executor(None, self.smtp.send_message, forwarded_msg)

    async def _get_subject_by_id(self, email_id: str) -> str:
        loop = asyncio.get_event_loop()
        _, messages = await loop.run_in_executor(None, self.imap.search, None, f'(HEADER Message-ID "{email_id}")')
        if messages[0]:
            msg_id = messages[0].split()[0]
            _, data = await loop.run_in_executor(None, self.imap.fetch, msg_id, '(BODY[HEADER.FIELDS (SUBJECT)])')
            header = data[0][1].decode()
            return header.replace("Subject: ", "").strip()
        return "Unknown"

    async def _fetch_raw_email_by_id(self, email_id: str) -> str:
        loop = asyncio.get_event_loop()
        _, messages = await loop.run_in_executor(None, self.imap.search, None, f'(HEADER Message-ID "{email_id}")')
        if messages[0]:
            msg_id = messages[0].split()[0]
            _, data = await loop.run_in_executor(None, self.imap.fetch, msg_id, '(RFC822)')
            return data[0][1].decode()
        return ""

        async def list_mailboxes(self):
            if not self.connected:
                raise Exception("Not connected")
            loop = asyncio.get_event_loop()
            status, mailboxes = await loop.run_in_executor(None, self.imap.list)
            if status == 'OK':
                logger.info("Available mailboxes:")
                for mbox in mailboxes:
                    logger.info(mbox)
            else:
                logger.error(f"Failed to list mailboxes: {mailboxes}")

# -----------------------------
# MCP Context & Server Setup
# -----------------------------

class EmailContext:
    def __init__(self, client: EmailClient):
        self.email_client = client


@asynccontextmanager
async def lifespan(app: Server) -> AsyncIterator[Dict]:
    email_client = EmailClient(EMAIL_CONFIG)
    await email_client.connect()
    ctx = EmailContext(client=email_client)
    yield {"email_ctx": ctx}
    await email_client.disconnect()


app = Server("mail-mcp-server", lifespan=lifespan)
sse = SseServerTransport("/messages/")


@app.list_tools()
async def list_tools() -> List[types.Tool]:
    ctx = app.request_context
    email_ctx = ctx.lifespan_context.get("email_ctx")
    if not email_ctx:
        raise ValueError("Failed to get email context")

    return [
        types.Tool(
            name="fetch_emails",
            description="Fetches emails from a specified folder (default: INBOX) up to a given limit.",
            inputSchema={
                "type": "object",
                "properties": {
                    "folder": {"type": "string", "default": "INBOX"},
                    "limit": {"type": "integer", "default": 10},
                    "search_criteria": {"type": "string", "default": "ALL"}
                },
                "required": ["folder", "limit", "search_criteria"]
            }
        ),
        types.Tool(
            name="reply_to_email",
            description="Replies to an email with a given ID.",
            inputSchema={
                "type": "object",
                "properties": {
                    "email_id": {"type": "string", "required": True},
                    "to": {"type": "string", "required": True},
                    "body": {"type": "string", "required": True}
                }
            }
        ),
        types.Tool(
            name="forward_email",
            description="Forwards an email to a new recipient.",
            inputSchema={
                "type": "object",
                "properties": {
                    "email_id": {"type": "string", "required": True},
                    "to": {"type": "string", "required": True},
                    "body": {"type": "string", "default": "Forwarded message"}
                }
            }
        )
    ]


@app.call_tool()
async def call_tool(name: str, arguments: dict) -> List[types.TextContent]:
    ctx = app.request_context
    email_ctx = ctx.lifespan_context.get("email_ctx")
    if not email_ctx:
        raise ValueError("Failed to get email context")

    client = email_ctx.email_client

    if name == "fetch_emails":
        folder = arguments.get("folder", "INBOX")
        limit = arguments.get("limit", 10)
        criteria = arguments.get("search_criteria", "ALL")
        emails = await client.fetch_emails(folder, limit, criteria)
        return [types.TextContent(type="text", text=json.dumps(emails))]

    elif name == "reply_to_email":
        email_id = arguments.get("email_id")
        to = arguments.get("to")
        body = arguments.get("body")
        await client.reply_to_email(email_id, to, body)
        return [types.TextContent(type="text", text="Email replied successfully")]

    elif name == "forward_email":
        email_id = arguments.get("email_id")
        to = arguments.get("to")
        body = arguments.get("body", "Forwarded message")
        await client.forward_email(email_id, to, body)
        return [types.TextContent(type="text", text="Email forwarded successfully")]

    raise ValueError(f"Unknown tool: {name}")


# -----------------------------
# SSE and App Routes
# -----------------------------

async def handle_sse(request):
    async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
        await app.run(streams[0], streams[1], app.create_initialization_options())


class AuthMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        if "/sse" in request.url.path or "/messages" in request.url.path:
            api_key = request.headers.get("api_key")
            # if not api_key:
            #     return JSONResponse({"error": "Missing authorization header"}, status_code=401)
        return await call_next(request)


middleware = [Middleware(AuthMiddleware)]

starlette_app = Starlette(
    debug=True,
    routes=[
        Route("/sse", endpoint=handle_sse),
        Mount("/messages/", app=sse.handle_post_message),
    ],
    middleware=middleware,
)

if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv()
    import uvicorn
    import os

    # Start the service using configurable environment variables
    uvicorn.run(
        "mail_mcp_server:starlette_app",  # Application module path
        host=os.getenv("SERVER_HOST", "0.0.0.0"),         # Host IP to bind (default: all interfaces)
        port=int(os.getenv("SERVER_PORT", "9383")),       # Port number (default: 9383)
        reload=(os.getenv("RELOAD", "False").lower() in ("true", "1", "yes")),  # Auto-reload flag
        log_level=os.getenv("LOG_LEVEL", "info")          # Logging verbosity (default: info)
    )